from typing import Tuple

import torch
import torch.nn as nn

from .feature_extractor import FeatureExtractor, FeatureExtractorV2
from .model_opts import ModelOpts
from .regressor import PoseRegressor, PoseRegressorDeploy, PoseRegressorSK
from .skeleton_encoder import SkeletonEncoder


class DYDModel(nn.Module):
    """
    V1:最初版本
    """

    def __init__(
            self,
            input_size: Tuple[int, int],
            model_opts: ModelOpts,
    ):
        super().__init__()
        self.feature_extractor = FeatureExtractor(input_size, model_opts)
        self.regressor = PoseRegressor(model_opts, True)

    def forward(
            self,
            images: torch.Tensor,
            extrinsic_matrices: torch.Tensor,
            intrinsic_matrices: torch.Tensor,
            # resize_ratios: torch.Tensor,
            masks: torch.Tensor,
    ):
        features = self.feature_extractor(images, extrinsic_matrices, intrinsic_matrices, masks)
        joint_angles, wrist_xfs, skel_scales, hand_prob = self.regressor(features)
        return joint_angles, wrist_xfs, skel_scales, hand_prob


class DYDModelV2(nn.Module):
    """
    训练版本
    """

    def __init__(
            self,
            input_size: Tuple[int, int],
            model_opts: ModelOpts,
    ):
        super().__init__()
        self.feature_extractor = FeatureExtractorV2(input_size, model_opts)
        self.regressor = PoseRegressor(model_opts, True)

        # 分类特征处理模块
        self.classification_head = nn.Sequential(
            # 额外的卷积层来提取分类特征
            nn.Conv2d(model_opts.nImageFeatureChannels, 512, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.BatchNorm2d(512),
            nn.AdaptiveAvgPool2d((1, 1)),  # 全局平均池化
            nn.Flatten(),
            nn.Dropout(0.3),
            nn.Linear(512, 256),
            nn.ReLU(inplace=True),
            nn.BatchNorm1d(256),
            nn.Dropout(0.2),
            nn.Linear(256, model_opts.numClasses)
        )

    def forward(
            self,
            images: torch.Tensor,
            masks: torch.Tensor,
            total_xfs_: torch.Tensor,
            world_to_ref_camera: torch.Tensor
    ):
        features = self.feature_extractor(images, masks, total_xfs_, world_to_ref_camera)
        joint_angles, wrist_xfs, skel_scales, hand_prob = self.regressor(features)
        return joint_angles, wrist_xfs, skel_scales, hand_prob

    def classify(
            self,
            images: torch.Tensor,
            masks: torch.Tensor,
            total_xfs_: torch.Tensor,
            world_to_ref_camera: torch.Tensor
    ):
        features = self.feature_extractor(images, masks, total_xfs_, world_to_ref_camera)
        class_logits = self.classification_head(features)

        return class_logits




class DYDModelV2Deploy(nn.Module):
    """
    部署版本
    """

    def __init__(
            self,
            input_size: Tuple[int, int],
            model_opts: ModelOpts,
    ):
        super().__init__()
        self.feature_extractor = FeatureExtractorV2(input_size, model_opts)
        self.regressor = PoseRegressorDeploy(model_opts, True)

    def forward(
            self,
            images: torch.Tensor,
            masks: torch.Tensor,
            total_xfs: torch.Tensor,
            world_to_ref_camera: torch.Tensor
    ):
        # images = images.reshape((1, 4, 1, 96, 96))
        features = self.feature_extractor(images, masks, total_xfs, world_to_ref_camera)
        joint_angles, rigid_pts, skel_scales, hand_side_logits = self.regressor(features)
        return joint_angles, rigid_pts, skel_scales, hand_side_logits


class DYDModelSK(nn.Module):
    """
    训练版本
    """

    def __init__(
            self,
            input_size: Tuple[int, int],
            model_opts: ModelOpts,
    ):
        super().__init__()
        self.feature_extractor = FeatureExtractorV2(input_size, model_opts)
        self.skeleton = SkeletonEncoder(
            [model_opts.nSkeletonFeatureChannels, *self.feature_extractor.backbone_out_feature_sizes])
        self.regressor = PoseRegressorSK(model_opts, True)

    def forward(
            self,
            images: torch.Tensor,
            masks: torch.Tensor,
            total_xfs_: torch.Tensor,
            joint_rotation_axes: torch.Tensor,
            joint_rest_positions: torch.Tensor,
            world_to_ref_camera: torch.Tensor
    ):
        features = self.feature_extractor.forward(images, masks, total_xfs_, world_to_ref_camera)

        skel_features = self.skeleton.forward(joint_rotation_axes, joint_rest_positions)
        img_skel_features = torch.cat([features, skel_features], dim=1)
        joint_angles, wrist_xfs, skel_scales, hand_prob = self.regressor(img_skel_features)
        return joint_angles, wrist_xfs, skel_scales, hand_prob
